import matplotlib.pyplot as plt
import numpy as np
import  tensorflow as tf


def plot_cost_accuracy(costs, train_accs, val_accs):
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
    # fig.suptitle("Lay dimension: [2, 3, 4, 3]")

    ax1.set_title("Loss")
    ax1.plot(costs)
    ax1.set_ylabel('Cost')
    ax1.set_xlabel('Epochs')

    ax2.set_title("Accuracy")
    ax2.plot(train_accs, label='Train')
    ax2.plot(val_accs, label='Validate')
    ax2.set_ylabel('Accuracy')
    ax2.set_xlabel('Epochs')
    ax2.legend()


def plot_decision_boundary(model, X, y):
    # Set min and max values and give it some padding
    x_min, x_max = X[0, :].min() - 1, X[0, :].max() + 1
    y_min, y_max = X[1, :].min() - 1, X[1, :].max() + 1
    h = 0.01
    # Generate a grid of points with distance h between them
    xx, yy = np.meshgrid(np.arange(x_min, x_max, h), np.arange(y_min, y_max, h))
    # Predict the function value for the whole grid
    Z = model(np.c_[xx.ravel(), yy.ravel()])
    if isinstance(Z, tf.Tensor):
        Z = tf.argmax(Z, axis=1).numpy()
    Z = Z.reshape(xx.shape)
    # Plot the contour and training examples
    # plt.title("Decision Boundary, Lay dimension: [2, 3, 4, 3]")
    plt.title("Decision Boundary")
    plt.contourf(xx, yy, Z, cmap=plt.cm.Spectral)
    plt.ylabel('x2')
    plt.xlabel('x1')
    plt.scatter(X[0, :], X[1, :], c=y.ravel(), cmap=plt.cm.Spectral)
